(*
  Copyright (c) 2009 Barry Schwartz

  Permission is hereby granted, free of charge, to any person
  obtaining a copy of this software and associated documentation
  files (the "Software"), to deal in the Software without
  restriction, including without limitation the rights to use,
  copy, modify, merge, publish, distribute, sublicense, and/or sell
  copies of the Software, and to permit persons to whom the
  Software is furnished to do so, subject to the following
  conditions:

  The above copyright notice and this permission notice shall be
  included in all copies or substantial portions of the Software.

  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
  EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
  OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
  NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
  HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
  WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
  FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
  OTHER DEALINGS IN THE SOFTWARE.
*)

module type Bits =
sig
  type t
  val of_int      : int -> t
  val to_int      : t -> int
  val nibblemask  : t -> int
  val nibbleshift : t -> t
  val nibbleblock : int
end

module type S =
sig
  type key
  type 'value t
  val empty    : 'value t
  val add      : key -> 'value -> 'value t -> 'value t
  val remove   : key -> 'value t -> 'value t
  val find     : key -> 'value t -> 'value (* Raises Not_found. *)
  val find_opt : key -> 'value t -> 'value option (* Returns None. *)
  val mem      : key -> 'value t -> bool
  val map      : ('value1 -> 'value2) -> 'value1 t -> 'value2 t
  val mapi     : (key -> 'value1 -> 'value2) -> 'value1 t -> 'value2 t
  val fold     : (key -> 'value -> 'x -> 'x) -> 'value t -> 'x -> 'x
  val iter     : (key -> 'value -> unit) -> 'value t -> unit
  val of_list  : (key * 'value) list -> 'value t
  val to_list  : 'value t -> (key * 'value) list
end

module Make(Key : Bits) =
struct
  type key = Key.t (* with sexp *)
  type 'value t = int array * 'value option array * key array

  let zero = Key.of_int 0

  let empty = ([| |], [| |], [| |])

  let add key value (keys, values, fullkeys) =

    let (keys, values, fullkeys) =
      if Array.length keys = 0 then
        (Array.make Key.nibbleblock (-1),
         Array.make Key.nibbleblock None,
         Array.make Key.nibbleblock zero)
      else
        (keys, values, fullkeys)
    in

    let rec try_block offset nibble shifted_key k v fk =
      if shifted_key = zero then
        begin
          v.(offset + nibble) <- Some value;
          fk.(offset + nibble) <- key;
          (k, v, fk);
        end
      else
        if k.(offset + nibble) < 0 then
          let new_offset = Array.length k in
            k.(offset + nibble) <- new_offset;
            let new_k = Array.append k (Array.make Key.nibbleblock (-1)) in
            let new_v = Array.append v (Array.make Key.nibbleblock None) in
            let new_fk = Array.append fk (Array.make Key.nibbleblock zero) in
            let new_nibble = Key.nibblemask shifted_key in
            let new_shifted_key = Key.nibbleshift shifted_key in
              try_block new_offset new_nibble new_shifted_key new_k new_v new_fk
        else
          let new_offset = k.(offset + nibble) in
          let new_nibble = Key.nibblemask shifted_key in
          let new_shifted_key = Key.nibbleshift shifted_key in
            try_block new_offset new_nibble new_shifted_key k v fk
    in
    let first_offset = 0 in
    let first_nibble = Key.nibblemask key in
    let first_shifted_key = Key.nibbleshift key in
      try_block first_offset first_nibble first_shifted_key keys values fullkeys

  let remove key (keys, values, fullkeys) =
    (* Note that this implementation of |remove| does not shrink the
       structure. *)
    
    if Array.length keys = 0 then
      (keys, values, fullkeys)
    else

      let rec try_block offset nibble shifted_key =
        if shifted_key = zero then
          match values.(offset + nibble) with
            | None -> (keys, values, fullkeys)
            | _ ->
                values.(offset + nibble) <- None;
                (keys, values, fullkeys)
        else
          let new_offset = keys.(offset + nibble) in
            if new_offset < 0 then
              (keys, values, fullkeys)
            else
              let new_nibble = Key.nibblemask shifted_key in
              let new_shifted_key = Key.nibbleshift shifted_key in
                try_block new_offset new_nibble new_shifted_key
      in
      let first_offset = 0 in
      let first_nibble = Key.nibblemask key in
      let first_shifted_key = Key.nibbleshift key in
        try_block first_offset first_nibble first_shifted_key

  let find_opt key (keys, values, _fullkeys) =

    if Array.length keys = 0 then
      None
    else

      let rec try_block offset nibble shifted_key =
        if shifted_key = zero then
          values.(offset + nibble)
        else
          let new_offset = keys.(offset + nibble) in
            if new_offset < 0 then
              None
            else
              let new_nibble = Key.nibblemask shifted_key in
              let new_shifted_key = Key.nibbleshift shifted_key in
                try_block new_offset new_nibble new_shifted_key
      in
      let first_offset = 0 in
      let first_nibble = Key.nibblemask key in
      let first_key = Key.nibbleshift key in
        try_block first_offset first_nibble first_key

  let find key trie =
    match find_opt key trie with
      | None -> raise Not_found
      | Some v -> v

  let mem key trie =
    match find_opt key trie with
      | None -> false
      | _    -> true

  let map f (keys, values, fullkeys) =
    let optional_f value =
      match value with
        | None -> None
        | Some v -> Some (f v)
    in
      (keys, Array.map optional_f values, fullkeys)

  let mapi f (keys, values, fullkeys) =
    let map_one i value =
      match value with
        | None -> None
        | Some v -> Some (f fullkeys.(i) v)
    in
      (keys, Array.mapi map_one values, fullkeys)

  let fold f (_keys, values, fullkeys) x =
    (* The order of folding is not specified, so maybe this function
       should be called |crumple| instead. :) *)
    let length = Array.length values in
    let rec loop i y =
      if i = length then
        y
      else
        match values.(i) with
          | None   -> loop (i + 1) y
          | Some v -> loop (i + 1) (f fullkeys.(i) v y)
    in
      loop 0 x

  let iter f (_keys, values, fullkeys) =
    let length = Array.length values in
      for i = 0 to length - 1 do
        match values.(i) with
          | None   -> ()
          | Some v -> f fullkeys.(i) v
      done

  let of_list p =
    List.fold_left (fun trie (k, v) -> add k v trie) empty p

  let to_list trie =
    let enlist_pair k v p =
      (k, v) :: p
    in
      fold enlist_pair trie []
end
